# Copyright (c) 2010
# The Regents of the University of Michigan
# All Rights Reserved

# Permission is granted to use, copy, create derivative works, and
# redistribute this software and such derivative works for any purpose,
# so long as the name of the University of Michigan is not used in
# any advertising or publicity pertaining to the use or distribution
# of this software without specific, written prior authorization. If
# the above copyright notice or any other identification of the
# University of Michigan is included in any copy of any portion of
# this software, then the disclaimer below must also be included.

# This software is provided as is, without representation or warranty
# of any kind either express or implied, including without limitation
# the implied warranties of merchantability, fitness for a particular
# purpose, or noninfringement.  The Regents of the University of
# Michigan shall not be liable for any damages, including special,
# indirect, incidental, or consequential damages, with respect to any
# claim arising out of or in connection with the use of the software,
# even if it has been or is hereafter advised of the possibility of
# such damages.

import hashlib
import os.path
import urlparse

from pydht import dht, lookup, routing_table


LOOKUP_CACHE_MAX = 50  # Arbitrary constant, the best kind. ;)


class EZDHTClient(lookup.LookupClient):
  '''Simplest reasonable client for the Vuze DHT.

  Synchronous, mapping-like interface. For now, does no caching and
  provides no storage of any kind. One consequence is that get()
  immediately after put()
  '''
  def __init__(self, address, app_prefix=''):
    '''Set up the EZDHTClient.

    address is required to be an (ip, port) tuple.

    app_prefix is automatically applied to all keys; recommended if
    someone else (perhaps another application of yours) might try to
    use the same keys for unrelated purposes.
    '''

    router = routing_table.RoutingTable(address)
    lookup.LookupClient.__init__(self, address, router)
    try:
      os.mkdir(os.path.expanduser('~/.ezdht'))
    except OSError:
      pass


    try:
      router.load(open(os.path.expanduser('~/.ezdht/routing_table'), 'rb'))
    except (routing_table.RoutingTableError, IOError):
      self.minibootstrap()
    self.version = 1
    self.lookup_cache = {}  # Map key -> [(node, spoof_id)]
    self.lookup_cache_used = set()

  def save(self):
    self.router.save(open(os.path.expanduser('~/.ezdht/routing_table'), 'wb'))

  def get(self, key):
    '''Return a list of values associated with the given key.

    Semantics are like the Vuze client\'s high-level interface and are
    what you want in the usual case: the key is an arbitrary-length
    byte sequence that is hashed with SHA-1. You get back a list of
    TransportValue objects (see DHT serialization module(s))
    corresponding to the given values. (In brief, it\'s an object
    containing the version, value, and originating IP of the value.)

    Comments in the Vuze source indicate that values are uniqified,
    although I have my doubts because that would imply that if two
    clients that used the same port, only one would be returned.

    If you\'re looking up a torrent, use its infohash as a raw string
    (i.e., if h is the hex string, use h.decode(\'hex\')). The value
    format is [ip_override:]tcp_port[;C][;udp_port], where everything
    in brackets is optional. Specifics can be gleaned from
    com.aelitis.azureus.plugins.tracker.dht.DHTTrackerPlugin in the
    Azureus/Vuze source code; the only non-obvious thing is that the ";C"
    denotes that BitTorrent protocol encryption is required.
    '''
    # Boy, lots of documentation for a one-line method!
    # TODO: implement Kademlia-recommended caching.
    return self.lookup_value(hashlib.sha1(key).digest())

  def put(self, key, value):
    '''Associate the given value (string) with the given key (string).

    If you want to store multiple values, right now you\'re stuck
    calling put() multiple times. There is no technical reason for this,
    it\'s just not implemented.

    The asymmetry with get() is necessary; several pieces of
    information (originator, version) make no sense to control on
    put() but are needed from get().
    '''
    hashed_key = hashlib.sha1(key).digest()
    nodes = self.lookup_cache.get(key)

    if nodes is None:
      nodes = self.lookup(hashed_key)
      self.lookup_cache[key] = nodes
      if len(self.lookup_cache) > LOOKUP_CACHE_MAX:
        for k in self.lookup_cache:
          if k not in self.lookup_cache_used:
            del self.lookup_cache[k]
            break
        self.lookup_cache_used = set()
    else:
      self.lookup_cache_used.add(key)

    version = self.version
    value = str(value)
    for node, spoof_id in nodes:
#      print 'storing it at %s' % node.dht_id.encode('hex')
      self.send_store_request(node, hashed_key, value, version, spoof_id)
    self.version += 1


def format_as_peers_list(values):
  '''Convert [TransportValue] -> [IP:port string].'''
  res = []
  for tv in values:
    if len(tv.bytes) <= 5 and tv.bytes.isdigit():
      port = tv.bytes
    else:
      port = dht.is_peer_record(tv.bytes)
    if port is not None:
      res.append('%s:%s' % (tv.originator[0], port))
  return res


def test(client):
  '''Test this module.'''
  import datetime
  print 'Using %s:%d' % client.address
  print 'Starts at: %s' % datetime.datetime.today()

  client.put('num', 1)
  print 'Put finished at: %s' % datetime.datetime.today()

  res = client.get('num')
  print 'Get finished at: %s' % datetime.datetime.today()
  print res

  client.put('num', 2)
  print 'Second put finished at: %s' % datetime.datetime.today()

  res = client.get('num')
  print 'Second get finished at: %s' % datetime.datetime.today()
  print res


def main(argv):
  import datetime
  import random
  myport = 4242
  ext_ip = argv[0]
  me = (ext_ip, myport)

  client = EZDHTClient(me)

  if len(argv) < 2:
    test(client)
  else:
    for key in argv[1:]:
      if key.startswith('magnet:'):
        parsed = urlparse.urlparse(key)
        if parsed.scheme == 'magnet':
          if not parsed.query:
            q = parsed.path[1:]  # Slice off the question mark.
          else:
            q = parsed.query
          qs = urlparse.parse_qs(q)
          xts = qs['xt']
          for xt in xts:
            if xt.startswith('urn:btih:'):
              key = xt[len('urn:btih:'):]

      key = key.decode('hex')
      for ip_port in format_as_peers_list(client.get(key)):
        print '%s:%d' % ip_port

  client.save()

if __name__ == '__main__':
  import sys
  main(sys.argv[1:])
